#!/usr/local/bin/python2.6
# Copyright (c) 2010-, WangFei, Cloud Team, CGCL,
# All rights reserved.
# This program configures multiple virtual machines(VM) as a MPI environment.
# Using format is as follow:
#    python setmpienv --machinefile machine
# Machine file includes list of IP and Root'Password of VM. 
# And the first line of machine will as the Master Node. 
# This program is only for specific virtual machine.
# 
#
import sys
import os
import pexpect
import time
import pp
import paramiko
import MySQLdb
import string
def expectandshow(child, expect):
	index=child.expect(expect)
	#print child.before,child.after,
	return index

def checkAllVM(ip,password):
	cmd="ssh " + ip	
   	scp_child=pexpect.spawn(cmd)
   	index=expectandshow(scp_child, \
                 ['Are you sure you want to continue connecting', \
                 "password:",pexpect.TIMEOUT, pexpect.EOF])
	if index == 0:
		scp_child.sendline('yes')
		index1=expectandshow(scp_child, ["password:",pexpect.TIMEOUT, pexpect.EOF])
		if index1 == 0:
			scp_child.sendcontrol('c')
			expectandshow(scp_child, [pexpect.EOF])
		else:
			expectandshow(scp_child, [pexpect.EOF])
		return 0
	elif index == 1:
		scp_child.sendcontrol('c')
		expectandshow(scp_child, [pexpect.EOF])
		return 0
	else: 
		return 1

def set_no_pwd(ip,username,password,ssh_port):
    
    #send_data_to_master(ip,username,password,ssh_port,"/root/.ssh/id_rsa.pub","/root/.ssh/id_rsa.pub")
    
    client = paramiko.SSHClient()
    client.load_system_host_keys()
    client.connect((ip), username=(username), password=(password), port=(ssh_port))
    sftpclient = client.open_sftp()
    sftpclient.put("/root/.ssh/id_rsa.pub","/root/.ssh/id_rsa.pub",callback=None)
    sftpclient.close()
    #client.close()

    cmd = "cat /root/.ssh/id_rsa.pub >> /root/.ssh/authorized_keys; chmod 600 /root/.ssh/authorized_keys"
    #client=paramiko.SSHClient()
    #client.load_system_host_keys()
    #client.connect((ip), username=(username), password=(password), port=(ssh_port))
    stdin,stdout,stderr = client.exec_command(cmd)
    print stdout.read()
    client.close()

def set_etc_hosts_no_pwd(host_address,username,all_host_list):
    cmd = "ssh "+username+"@" + host_address + " \"echo \'127.0.0.1       localhost.localdomain localhost\' > /etc/hosts;"
    for aip in all_host_list:
        cmd += "echo \'"+aip+" "+aip+"\' >> /etc/hosts;"
    cmd += "/bin/hostname "+host_address +";"
    cmd += " echo '' > /etc/hosts.equiv;"
    for ip1 in all_host_list:
        cmd += "echo \'" +ip1 + "\' >> /etc/hosts.equiv;"
    cmd += "echo '' > /root/.rhosts;"
    for ip2 in all_host_list:
        cmd += "echo \'" +ip2 + "\' >> /root/.rhosts;"
    HOSTNAME=host_address
    NETWORKING="yes"
    NETWORKING_IPV6="no"
    cmd += " echo \'" + "NETWORKING="+NETWORKING+"\' > /etc/sysconfig/network;" \
            + "echo \'NETWORKING_IPV6=" + NETWORKING_IPV6 + "\' >> /etc/sysconfig/network ;" \
            + "echo \'HOSTNAME=" + HOSTNAME + "\' >> /etc/sysconfig/network;\" "
    #print cmd
    out = os.popen(cmd)

def set_salve_nfs_rsh_no_pwd(Salve_IP,Master_IP,username):
    cmd = "ssh "+username+"@" + Salve_IP + " \"sed -i \'7d\' /etc/fstab;" + \
         "echo \'"+ Master_IP+":/home /home nfs defaults 0 0\' >> /etc/fstab ;" + \
         "/sbin/chkconfig netfs on;" + \
                "mount /home;" + \
                "/sbin/chkconfig rsh on;" + \
                "/etc/init.d/xinetd start;\""
    out = os.popen(cmd)
    print out.read()

def set_master_rsh_no_pwd(Master_IP,username):
    cmd ="ssh "+username+"@" + Master_IP + " \"/sbin/chkconfig rsh on;" + \
              "/etc/init.d/xinetd start;\""
    out = os.popen(cmd)
    print out.read()    

def set_master_rsh(Master_IP,username,password,ssh_port):
    cmd = "/sbin/chkconfig rsh on;" + \
              "/etc/init.d/xinetd start;"
    client = paramiko.SSHClient()
    client.load_system_host_keys()
    client.connect((Master_IP), username=(username), password=(password), port=(ssh_port))
    stdin,stdout,stderr = client.exec_command(cmd)
    print stdout.read()
    client.close()

def set_master_nfs_no_pwd(Master_IP,username):
    cmd = "ssh "+username+"@" + Master_IP + " \"/sbin/chkconfig nfs on;" + \
        	"/sbin/chkconfig nfslock on;" + \
        	"/etc/init.d/nfslock start;" + \
        	"/etc/init.d/nfs start;\""
    out = os.popen(cmd)
    print out.read()

def set_master_nfs(Master_IP,username,password,ssh_port):
    cmd = "/sbin/chkconfig nfs on;" + \
            "/sbin/chkconfig nfslock on;" + \
            "/etc/init.d/nfslock start;" + \
            "/etc/init.d/nfs start"
    client = paramiko.SSHClient()
    client.load_system_host_keys()
    client.connect((Master_IP), username=(username), password=(password), port=(ssh_port))
    stdin,stdout,stderr = client.exec_command(cmd)
    print stdout.read()
    client.close()	

def set_mpd_port_no_pwd(Master_IP, total_ip, username, all_host_list):
    cmd = "ssh "+username+"@" + Master_IP + " \""
    cmd += "echo '' > /root/mpd.hosts;"
    for ip in all_host_list:
        cmd += "echo "+ip+" >> /root/mpd.hosts;"
    cmd += "echo '' > /home/cgcl/all_vm_ip_list;"
    for ip in all_host_list:
        cmd += "echo "+ip+" >> /home/cgcl/all_vm_ip_list;"
    cmd += "mpdallexit; mpdboot --rsh=rsh -n "+str(total_ip)+" --file=/root/mpd.hosts; \""
    out = os.popen(cmd)
    print out.read()

def set_mpd_port(Master_IP, total_ip, username, password, ssh_port, all_host_list):
    cmd = "echo '' > /root/mpd.hosts;"
    for ip in all_host_list:
        cmd += "echo "+ip+" >> /root/mpd.hosts;"
    cmd += "echo '' > /home/cgcl/all_vm_ip_list;"
    for ip in all_host_list:
        cmd += "echo "+ip+" >> /home/cgcl/all_vm_ip_list;"
    cmd += "mpdallexit; mpdboot --rsh=rsh -n "+str(total_ip)+" --file=/root/mpd.hosts"
    client = paramiko.SSHClient()
    client.load_system_host_keys()
    client.connect((Master_IP), username=(username), password=(password), port=(ssh_port))
    stdin,stdout,stderr = client.exec_command(cmd)
    print stdout.read()
    client.close()

def GetIpAndPassword(job_id,conn):
    global all_host_list
    global total_ip
    all_host_list=[]
    total_ip = 0
    cnt = conn.execute("select vm_ip from env_resource, vm_info where env_id="+job_id+" and env_resource.vm_id=vm_info.vm_id and vm_info.vm_type='master_node'")
    if cnt != 1:
        print "Not Master Node,Fail!" 
        return 0
    rows=conn.fetchall()
    all_host_list.append(rows[0][0])
    total_ip += 1
    
    conn.execute("select vm_ip from env_resource, vm_info where env_id="+job_id+" and env_resource.vm_id=vm_info.vm_id and vm_info.vm_type='salve_node'")
    
    rows=conn.fetchall()
    for temp in rows:
        all_host_list.append(temp[0])
        total_ip += 1
    #print all_host_list
    #print password_list
    return 1
    
def testSSHtoVM():
	global all_host_list
	global password_list
	global total_ip
	#start_time_total = time.time()
	#start_time = time.time()
	ppservers = ()
	job_server = pp.Server(total_ip, ppservers=ppservers)
	#print "Starting pp with", job_server.get_ncpus(), "workers"
	jobs = []
	index = 0
	while index < total_ip:
		host_address=all_host_list[index]
		password = password_list[index]
		jobs.append(job_server.submit(checkAllVM, (host_address,password,), (expectandshow,),("pexpect",)))
		index += 1
	for job in jobs:
        	job()
	#print "Time elapsed: ", time.time() - start_time, "s"
    #job_server.print_stats()


def setSSH():
    global all_host_list
    global total_ip
    global password_list
    ssh_port=22
    username="root"
    ppservers = ()
    job_server = pp.Server(total_ip, ppservers=ppservers)
    #print "Starting pp with", job_server.get_ncpus(), "workers"
    jobs = []
    index = 0
    while index < total_ip:
        host_address=all_host_list[index]
        password = password_list[index]
        jobs.append(job_server.submit(set_no_pwd, (host_address,username,password,ssh_port,), (),("paramiko",)))
        index += 1

    for job in jobs:
        job()


def set_mpi_env_no_pwd():
    global all_host_list
    username="root"
    Master_IP=all_host_list[0]
    #start_time_total = time.time()
    #start_time = time.time()
    ppservers = ()
    job_server = pp.Server(total_ip, ppservers=ppservers)
    #print "Starting pp with", job_server.get_ncpus(), "workers"
    jobs = []
    index = 0
    while index < total_ip:
        host_address=all_host_list[index]
        password = password_list[index]
        jobs.append(job_server.submit(set_etc_hosts_no_pwd, (host_address,username,all_host_list,), (),("os",)))
        index += 1

    for job in jobs:
        job()

    #print "Time elapsed: ", time.time() - start_time, "s"
    #job_server.print_stats()
    #start_time = time.time()
    ppservers = ()
    job_server = pp.Server(2, ppservers=ppservers)
    #print "Starting pp with", job_server.get_ncpus(), "workers"

    jobs4 = []
    jobs4.append(job_server.submit(set_master_nfs_no_pwd, (Master_IP,username,), (), ("os",)))
    jobs4.append(job_server.submit(set_master_rsh_no_pwd, (Master_IP,username,), (), ("os",)))
    for job in jobs4:
        job()
    #print "Time elapsed: ", time.time() - start_time, "s"
    #job_server.print_stats()
    #start_time = time.time()
    ppservers = ()
    job_server = pp.Server(total_ip-1, ppservers=ppservers)
    #print "Starting pp with", job_server.get_ncpus(), "workers"

    jobs1 = []
    index = 1
    while index < total_ip:
        host_address=all_host_list[index]
        jobs1.append(job_server.submit(set_salve_nfs_rsh_no_pwd, (host_address,Master_IP,username,), (), ("os",)))
        index += 1
    for job in jobs1:
        job()
    #print "Time elapsed: ", time.time() - start_time, "s"
    #job_server.print_stats()
    #print "Master ip: "
    #print "*"
    #print Master_IP
    #print "*"

    set_mpd_port_no_pwd(Master_IP,total_ip,username,all_host_list)
    #print "Time elapsed: ", time.time() - start_time_total, "s"

def set_mpi_env():
    global all_host_list
    global total_ip
    global password_list
    ssh_port=22
    username="root"
    Master_IP=all_host_list[0]    
    #start_time_total = time.time()
    #start_time = time.time()
    ppservers = ()
    job_server = pp.Server(total_ip, ppservers=ppservers)
    #print "Starting pp with", job_server.get_ncpus(), "workers"
    jobs = []
    index = 0
    while index < total_ip:
        host_address=all_host_list[index]
        password = password_list[index]
        jobs.append(job_server.submit(set_etc_hosts, (host_address,username,password,ssh_port,all_host_list,), (),("paramiko",)))
        index += 1

    for job in jobs:
        job()
    
    #print "Time elapsed: ", time.time() - start_time, "s"
    #job_server.print_stats()
    #start_time = time.time()
    ppservers = ()
    job_server = pp.Server(2, ppservers=ppservers)
    #print "Starting pp with", job_server.get_ncpus(), "workers"
    
    password=password_list[0]
    jobs4 = []
    jobs4.append(job_server.submit(set_master_nfs, (Master_IP,username,password,ssh_port,), (), ("paramiko",)))
    jobs4.append(job_server.submit(set_master_rsh, (Master_IP,username,password,ssh_port,), (), ("paramiko",)))
    for job in jobs4:
        job()
    #print "Time elapsed: ", time.time() - start_time, "s"
    #job_server.print_stats()
    #start_time = time.time()
    ppservers = ()
    job_server = pp.Server(total_ip-1, ppservers=ppservers)
    #print "Starting pp with", job_server.get_ncpus(), "workers"

    jobs1 = []
    index = 1
    while index < total_ip:
        host_address=all_host_list[index]
        password=password_list[index]
        jobs1.append(job_server.submit(set_salve_nfs_rsh, (host_address,Master_IP,username,password,ssh_port,), (), ("paramiko",)))
        index += 1 
    for job in jobs1:
        job()
    #print "Time elapsed: ", time.time() - start_time, "s"
    #job_server.print_stats()
    #print "Master ip: "
    #print "*"
    #print Master_IP
    #print "*"

    password=password_list[0]
    set_mpd_port(Master_IP,total_ip,username,password,ssh_port,all_host_list)
    #print "Time elapsed: ", time.time() - start_time_total, "s"


def isDigtal(myNum):
    nums = string.digits
    if type(myNum) is not str:
        return False
    else:
        for temp in myNum:
            if temp not in nums:
                return False
        return True

def checkPara( conn ):
    global job_id
    if not conn:
        print "Can not connect to database, Fail!"
        return -1
    
    if len(sys.argv) != 2:
        return 0
    if not isDigtal(sys.argv[1]):
        return 0
    job_id = str(sys.argv[1])
    cnt = conn.execute("select env_id from mpi_env where status='running' and env_id="+job_id)    
    if cnt != 1:
        print "Not job ",job_id
        return 0
    return 1

all_host_list=[]
password_list=[]
total_ip=0
job_id="0"

str_command = "Command format: ./setmpienv job_id"

if __name__ == "__main__":
    connect=MySQLdb.connect(host="localhost",port=3306,user="cloud",passwd="cloudpass",db="cloud")
    conn=connect.cursor()
    
    start_time = time.time()
    flag = checkPara(conn)
    if flag == -1:
        sys.exit(0)
    if flag == 0:
        print str_command
        sys.exit(0)
    
    flag = GetIpAndPassword(job_id,conn)
    if flag == 0:
        sys.exit(0)
    #start_time=time.time()
    #print "tessSSHtoVM:"
    #testSSHtoVM()
    #print "testSSHtoVm time: ", time.time() - start_time, "s"
    #print "setSSH:"
    #start_time=time.time()
    #setSSH()
    #print "setSSH time: ", time.time() - start_time, "s"
    #set_mpi_env()
    
    start_time=time.time()
    set_mpi_env_no_pwd()
    print "setMpiEnv time: ", time.time() - start_time, "s"
    print "\n\nIt's successful."
    conn.close()
    connect.close()

